import cv2
import torch
import matplotlib.pyplot as plt
import numpy as np
import json
import subprocess
import re
from dataclasses import dataclass
import os


@dataclass
class PostProcessParams:
    """
    kotlin代码中处理heatmap的相关参数

    :param width: 原图的宽度
    :param height: 原图的高度
    :param heatValueThreshold: 第一步对热图进行过滤的阈值
    :param distThreshold: 非极大抑制时，两个关键点之间的距离阈值
    :param maxDistanceDifference: 对nms后得到的关键点使用匹配算法时的距离阈值
    :param angleDifferenceThreshold: 对nms后得到的关键点使用匹配算法时的角度阈值
    """

    width: int = 512
    height: int = 384
    heatValueThreshold: float = 0.5
    distThreshold: float = 10.0

    maxDistanceDifference: float = 30
    angleDifferenceThreshold: float = round(np.pi / 4, 2)


def detect_convert_heatmaps_to_keypoints(heatmaps, postprocess_params):
    """
    将热图转换为关键点，并区分Code128和QR码的关键点。
    :param heatmaps: 一个batch的热图，shape为(batch_size, 3, height, width)
    :return: 一个batch的关键点，每张图片返回关键点集合
    """
    device = heatmaps.device
    batch_size, _, height, width = heatmaps.shape
    keypoints = []

    for i in range(batch_size):
        # 二维位置热图
        position_heatmap = heatmaps[i, 0].clone()
        # 获取热图中大于阈值的点的坐标
        y_coords, x_coords = torch.where(
            position_heatmap > postprocess_params.heatValueThreshold
        )
        if len(y_coords) > 0:
            values = position_heatmap[y_coords, x_coords]
            # 按照热值的大小对坐标进行排序
            top_k_indices = torch.topk(values, k=len(values), largest=True).indices
            top_k_coords = torch.stack(
                (x_coords[top_k_indices], y_coords[top_k_indices]), dim=1
            )

            for coord in top_k_coords:
                if all(
                    torch.norm(
                        coord.float() - torch.tensor(kp[:2], device=device).float(),
                        p=2,
                    )
                    > postprocess_params.dist_threshold
                    for kp in keypoints
                ):
                    # sin_val = heatmaps[i, 1][coord[1], coord[0]]
                    # cos_val = heatmaps[i, 2][coord[1], coord[0]]

                    sin_val = torch.mean(
                        heatmaps[i, 1][
                            coord[1],
                            coord[0],
                        ]
                    )
                    cos_val = torch.mean(
                        heatmaps[i, 2][
                            coord[1],
                            coord[0],
                        ]
                    )

                    # 计算角度值
                    angle_rad = np.arctan2(sin_val, cos_val)
                    keypoints.append([coord[0].item(), coord[1].item(), angle_rad])
    return keypoints


def run_kotlin_script(heatmap, postprocess_params):
    heatmap_json = json.dumps(heatmap.tolist()).encode("utf-8")
    command = [
        r"D:\AndroidStudio\jbr\bin\java.exe",
        r"-javaagent:D:\AndroidStudio\lib\idea_rt.jar=56231:D:\AndroidStudio\bin",
        "-Dfile.encoding=UTF-8",
        "-classpath",
        r"C:\Users\Administrator\Desktop\aidc_kotelin\common\build\classes\kotlin\main;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\multik-kotlin-jvm\0.2.3\742edf7899ccbdd2b76b01fcd4e0a2c69f30fa9d\multik-kotlin-jvm-0.2.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.fasterxml.jackson.module\jackson-module-kotlin\2.13.3\5e1944bc4efe2f5a5dcf286a83976fed61c1c65b\jackson-module-kotlin-2.13.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlin\kotlin-stdlib\1.9.22\d6c44cd08d8f3f9bece8101216dbe6553365c6e3\kotlin-stdlib-1.9.22.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\junit\junit\4.12\2973d150c0dc1fefe998f834810d68f278ea58ec\junit-4.12.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.google.zxing\javase\3.4.1\4eddd2ff030c0d4f368d3cc00131ff6a58bc88e0\javase-3.4.1.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.google.zxing\core\3.4.1\1869da97e9b2b60b5ff2fcaf55899174b93ae25d\core-3.4.1.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.fasterxml.jackson.core\jackson-databind\2.13.3\56deb9ea2c93a7a556b3afbedd616d342963464e\jackson-databind-2.13.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.fasterxml.jackson.core\jackson-annotations\2.13.3\7198b3aac15285a49e218e08441c5f70af00fc51\jackson-annotations-2.13.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\multik-default-jvm\0.2.3\767e0932c58dcd2bfa3f0d2535df9a3053a52046\multik-default-jvm-0.2.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\multik-core-jvm\0.2.3\7ea1270b703a02b7d8b978a350aa6d5df03987e3\multik-core-jvm-0.2.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\kotlinx-serialization-json-jvm\1.0.1\900f8d34f091c8a957eca49e0cf6bf7d7006482d\kotlinx-serialization-json-jvm-1.0.1.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\kotlinx-coroutines-core-jvm\1.6.0\f3b8fd26c2e76d2f18cbc36aacb6e349fcb9fd5f\kotlinx-coroutines-core-jvm-1.6.0.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains\annotations\13.0\919f0dfe192fb4e063e7dacadee7f8bb9a2672a9\annotations-13.0.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.hamcrest\hamcrest-core\1.3\42a25dc3219429f0e5d060061f71acb49bf010a0\hamcrest-core-1.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.beust\jcommander\1.78\a3927de9bd6f351429bcf763712c9890629d8f51\jcommander-1.78.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.fasterxml.jackson.core\jackson-core\2.13.3\a27014716e4421684416e5fa83d896ddb87002da\jackson-core-2.13.3.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlin\kotlin-stdlib-jdk8\1.9.0\e000bd084353d84c9e888f6fb341dc1f5b79d948\kotlin-stdlib-jdk8-1.9.0.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\kotlinx-serialization-core-jvm\1.0.1\d0c5ab89372940ed3fb77952e79b9cd5d466d86\kotlinx-serialization-core-jvm-1.0.1.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlin\kotlin-stdlib-jdk7\1.9.0\f320478990d05e0cfaadd74f9619fd6027adbf37\kotlin-stdlib-jdk7-1.9.0.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlin\kotlin-reflect\1.9.22\761ab33cf85f03a3a24595e1b92639d82e7a9564\kotlin-reflect-1.9.22.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\com.github.jai-imageio\jai-imageio-core\1.4.0\fb6d79b929556362a241b2f65a04e538062f0077\jai-imageio-core-1.4.0.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.bio\npy\0.3.5\d6c84e30c331136b1814978276b099086d56de34\npy-0.3.5.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.apache.commons\commons-csv\1.10.0\8669bee353424c3223c93723291b5c3753260c1c\commons-csv-1.10.0.jar;C:\Users\Administrator\.gradle\caches\modules-2\files-2.1\org.jetbrains.kotlinx\multik-openblas-jvm\0.2.3\d431f401d3fa3cbc58c1f5769c71413b20bfbb7e\multik-openblas-jvm-0.2.3.jar",
        "com.action.ai.common.keypointdetector.KeypointDetectorPostprocessKt",
    ]

    # 添加要传递给Kotlin程序的参数
    # kotlin_args_data = json.dumps(np.zeros([10, 10]).tolist()).encode("utf-8")
    # 传入kotlin图像的width和height参数,和float类型的阈值参数的字符形式
    kotlin_args = [
        str(postprocess_params.width),
        str(postprocess_params.height),
        str(postprocess_params.heatValueThreshold),
        str(postprocess_params.distThreshold),
        str(postprocess_params.maxDistanceDifference),
        str(postprocess_params.angleDifferenceThreshold),
    ]
    # 将参数添加到命令列表
    command.extend(kotlin_args)
    output = subprocess.run(
        command,
        input=heatmap_json,
        shell=True,
        capture_output=True,
    )

    output_stdout = output.stdout.decode("utf-8")
    return output_stdout


def detect_image(
    image_path, model, postprocess_params, input_tensor_height, input_tensor_width
):
    # 读取图像
    images_rgb = cv2.imread(image_path)
    images_rgb = cv2.resize(images_rgb, (input_tensor_width, input_tensor_height))
    images_rgb = cv2.cvtColor(images_rgb, cv2.COLOR_BGR2RGB)
    images_gray = cv2.cvtColor(images_rgb, cv2.COLOR_BGR2GRAY)
    images_gray = images_gray / 255.0
    images_tensor = (
        torch.tensor(images_gray)
        .reshape(1, 1, input_tensor_height, input_tensor_width)
        .float()
    )
    outputs = model(images_tensor)
    outputs = outputs.detach()
    outputs_np = outputs.numpy()

    # with open("test.json", "w") as f:
    #     json.dump(outputs_np.tolist(), f)
    # 调用kotlin程序,传入神经网络输出的热图，得到关键点坐标的字符串。字符串为kotlin程序的标准输出。
    predicted_keypoints_str = run_kotlin_script(outputs_np, postprocess_params)

    # 若是字符中含有关键点信息，则使用正则表达式提取关键点坐标并绘制在图像上
    if "WHITE_LABEL=" in predicted_keypoints_str:
        keypoints_inf = predicted_keypoints_str.split("separator")[0]
        matched_keypoints_inf = predicted_keypoints_str.split("separator")[1]

        keypoints = re.findall(r"-?\d+\.\d+", keypoints_inf)
        keypoints = list(map(float, keypoints))
        keypoints_array = np.array(keypoints).reshape(-1, 1, 4).astype(np.float32)

        matched_keypoints = re.findall(r"-?\d+\.\d+", matched_keypoints_inf)
        matched_keypoints = list(map(float, matched_keypoints))
        matched_keypoints_array = (
            np.array(matched_keypoints).reshape(-1, 4, 2).astype(np.int32)
        )

        cv2.polylines(
            images_rgb,
            matched_keypoints_array,
            isClosed=True,
            color=(0, 255, 0),
            thickness=1,
        )
        # 定义颜色
        color_keypoints = (255, 0, 0)  # 红色
        color_arrow = (0, 0, 255)  # 蓝色
        arrow_length = 12  # 箭头长度

        # 遍历所有检测到的关键点
        for kp in keypoints_array:
            # 调整关键点坐标考虑填充和缩放的影响
            kp = kp.squeeze()
            x = kp[0]
            y = kp[1]
            # 绘制代表角度的箭头
            end_x = x + arrow_length * kp[2]
            end_y = y + arrow_length * kp[3]
            # 绘制关键点
            cv2.arrowedLine(
                images_rgb,
                (int(x), int(y)),
                (int(end_x), int(end_y)),
                color_arrow,
                2,
                tipLength=0.35,
            )
            cv2.circle(images_rgb, (int(x), int(y)), 3, color_keypoints, -1)
    else:
        print("not matched")

    # keypoints = detect_convert_heatmaps_to_keypoints(outputs, score_threshold, 10)

    return images_rgb, outputs_np
